import os
import requests
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
from whoosh.index import create_in, open_dir, exists_in
from whoosh.fields import Schema, TEXT
from whoosh.qparser import QueryParser, FuzzyTermPlugin
from whoosh.analysis import StemmingAnalyzer
from openai import OpenAI
import PyPDF2
from docx import Document
import pandas as pd

# 初始化 Flask
app = Flask(__name__, static_folder='dist', static_url_path='')
CORS(app)

# 初始化 DeepSeek 客户端
client = OpenAI(api_key="sk-0bb2e1837ff7453c873c259368dbdf8f", base_url="https://api.deepseek.com")
dbaoclient = OpenAI(
    # 此为默认路径，您可根据业务所在地域进行配置
    base_url="https://ark.cn-beijing.volces.com/api/v3",
    # 从环境变量中获取您的 API Key。此为默认方式，您可根据需要进行修改
    api_key="111a646b-6523-4b91-9f81-246fa4ddcad5",
)

# 定义知识库文件夹路径
KNOWLEDGE_BASE_FOLDER = r"E:\\BijiApp\\web_app\\python_part\\Demo1\\TEMP"

# 定义 Whoosh 索引路径
INDEX_DIR = "whoosh_index"

# 定义支持模糊匹配的分析器
analyzer = StemmingAnalyzer()

# 用户对话历史（以用户 ID 为键）
user_sessions = {}

chat_mode_type  = 1 # 1 豆包 2 deepseek

def create_search_index(folder_path, index_dir):
    """创建或更新 Whoosh 倒排索引，支持模糊匹配"""
    schema = Schema(content=TEXT(stored=True, analyzer=analyzer))
    if not os.path.exists(index_dir):
        os.mkdir(index_dir)
    index = create_in(index_dir, schema)
    writer = index.writer()

    for filename in os.listdir(folder_path):
        file_path = os.path.join(folder_path, filename)
        content = ""
        print(f"Processing file: {filename}")

        if filename.endswith('.txt'):
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
        elif filename.endswith('.pdf'):
            with open(file_path, 'rb') as file:
                reader = PyPDF2.PdfReader(file)
                for page in reader.pages:
                    content += page.extract_text()
        elif filename.endswith('.docx'):
            doc = Document(file_path)
            for paragraph in doc.paragraphs:
                content += paragraph.text + "\n"
        elif filename.endswith('.xlsx'):
            df = pd.read_excel(file_path)
            content = df.to_string(index=False)

        print(f"Content: {content[:100]}...")  # 打印前 100 个字符
        if content:
            writer.add_document(content=content)
    writer.commit()
    print("Index updated successfully.")

def search_index(query, index_dir, top_k=3):
    """在 Whoosh 索引中检索相关文档，支持模糊匹配"""
    index = open_dir(index_dir)
    with index.searcher() as searcher:
        query_parser = QueryParser("content", index.schema)
        query_parser.add_plugin(FuzzyTermPlugin())
        query = query_parser.parse(query + "~2")  # 允许最多 2 个字符的编辑距离
        print(f"Parsed Query: {query}")  # 打印解析后的查询
        results = searcher.search(query, limit=top_k)
        #print(f"Search Results: {results}")  # 打印检索结果
        return [result['content'] for result in results]

def search_web(query):
    """联网搜索（使用博查AI开放平台的Web Search API）"""
    api_key = "sk-8f39cbc1c00b4fefb19c335f965d7751"  # 替换为你的博查AI API Key
    url = "https://api.bochaai.com/v1/web-search"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "query": query,
        "summary": True,
        "freshness": "noLimit",
        "count": 50
    }
    try:
        print("get boshi ")
        response = requests.post(url, headers=headers, json=data)
        if response.status_code == 200:
            results = response.json()
            #print("get boshi "+str(results.get('data')))
            snippets = []
            # 检查 data 字段是否存在且为列表类型
            if isinstance(results.get('data').get('webPages').get('value'), list):
                #print("get boshi data ")
                for item in results['data']['webPages']['value']:
                    if isinstance(item, dict):
                        # 根据实际返回结构调整字段名
                        snippet = item.get('summary', '')
                        #print("get boshi snippet")
                        if snippet:
                            snippets.append(snippet)
            return " ".join(snippets)
        else:
            print(f"Failed to fetch search results. Status code: {response.status_code}, Response: {response.text}")
            return None
    except requests.RequestException as e:
        print(f"Request error: {e}")
        return None

def call_deepseek_api(messages):
    """通过 OpenAI SDK 调用 DeepSeek API"""
    try:
        if chat_mode_type == 2 :
            print("call_deepseek_api ....")
            response = client.chat.completions.create(
                model="deepseek-chat",
                messages=messages
            )
        else :
            print("call_doubao_api ....")
            response = dbaoclient.chat.completions.create(
                # 指定您创建的方舟推理接入点 ID，此处已帮您修改为您的推理接入点 ID
                model="doubao-1-5-vision-pro-32k-250115",
                messages=messages,
            )
        #print("Request Data:", messages)
        print("Response:", response)
        return response.choices[0].message.content
    except Exception as e:
        print("Error:", e)
        return None

# 创建 Whoosh 索引（如果索引不存在）
if not os.path.exists(INDEX_DIR):
    print("Creating search index...")
    create_search_index(KNOWLEDGE_BASE_FOLDER, INDEX_DIR)

@app.route('/chat', methods=['POST'])
def chat():
    """支持多轮对话的接口，同时支持本地检索和联网搜索"""
    user_id = request.json.get('user_id')  # 用户唯一标识
    user_query = request.json.get('query')  # 用户输入
    search_type = request.json.get('search_type', 'local')  # 搜索类型：local（本地）或 web（联网）

    # 获取用户的对话历史
    if user_id not in user_sessions:
        user_sessions[user_id] = []

    # 将用户输入添加到对话历史
    if chat_mode_type == 2 : 
        user_sessions[user_id].append({"role": "user", "content": user_query}) # deepseek 格式
    else :
        user_sessions[user_id].append({"role": "user", "content": [{"type": "text", "text": user_query}]}) # 豆包格式

    # 根据搜索类型获取上下文
    if search_type == 'local':
        # 本地检索
        relevant_docs = search_index(user_query, INDEX_DIR, top_k=3)
        context = "\n".join(relevant_docs)
    elif search_type == 'web':
        # 联网搜索
        context = search_web(user_query)
    else:
        return jsonify({"status": "error", "message": "Invalid search type."})
    #print("context:"+context)

    # 将上下文添加到对话历史
    if context:
        if chat_mode_type == 2 : 
            user_sessions[user_id].append({"role": "user", "content": f"上下文：{context}"}) # deepseek 格式
        else :
            user_sessions[user_id].append({"role": "user", "content": [{"type": "text", "text": str(context)}]}) # 豆包格式
        

    # 调用 DeepSeek API
    assistant_reply = call_deepseek_api(user_sessions[user_id])
    print("assistant_reply")
    if assistant_reply:
        # 将助手回复添加到对话历史
        user_sessions[user_id].append({"role": "assistant", "content": assistant_reply})
        return jsonify({
            "user_id": user_id,
            "response": assistant_reply
        })
    else:
        return jsonify({"status": "error", "message": "Failed to generate response."})

@app.route('/update_index', methods=['POST'])
def update_index():
    """更新 Whoosh 索引"""
    try:
        create_search_index(KNOWLEDGE_BASE_FOLDER, INDEX_DIR)
        return jsonify({"status": "success", "message": "Index updated successfully."})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)})


@app.after_request
def add_header(response):
    response.headers['Cache-Control'] = 'no-store, must-revalidate'
    response.headers['Pragma'] = 'no-cache'
    response.headers['Expires'] = '0'
    return response

@app.route('/')
def index():
    return send_from_directory('dist', 'index.html')

if __name__ == '__main__':
    app.run(port=5000)